import tkinter as tk
from tkinter import ttk, filedialog, messagebox
from tkinterdnd2 import DND_FILES
from image_processor import ImageProcessor
from docx import Document
import os

class OCRApp:
    def __init__(self, root):
        self.root = root
        self.root.title("图像文字识别转Word文档工具")
        self.root.geometry("800x600")
        
        # 配置根窗口的网格权重，使其可以跟随窗口调整大小
        self.root.grid_rowconfigure(0, weight=1)
        self.root.grid_columnconfigure(0, weight=1)
        
        # 存储选择的图片路径
        self.image_paths = []
        
        self.create_widgets()
        
    def create_widgets(self):
        """创建GUI界面组件"""
        # 创建主框架
        main_frame = ttk.Frame(self.root, padding="10")
        main_frame.grid(row=0, column=0, sticky=(tk.W, tk.E, tk.N, tk.S))
        
        # 配置主框架的网格权重
        main_frame.grid_rowconfigure(0, weight=1)
        main_frame.grid_columnconfigure(0, weight=1)
        main_frame.grid_columnconfigure(1, weight=1)
        
        # 图片列表框
        list_frame = ttk.LabelFrame(main_frame, text="待处理图片列表", padding="5")
        list_frame.grid(row=0, column=0, columnspan=2, sticky=(tk.W, tk.E, tk.N, tk.S))
        
        # 配置列表框的网格权重
        list_frame.grid_rowconfigure(0, weight=1)
        list_frame.grid_columnconfigure(0, weight=1)
        
        self.listbox = tk.Listbox(list_frame, width=70, height=15)
        self.listbox.grid(row=0, column=0, sticky=(tk.W, tk.E, tk.N, tk.S))
        
        scrollbar = ttk.Scrollbar(list_frame, orient=tk.VERTICAL, command=self.listbox.yview)
        scrollbar.grid(row=0, column=1, sticky=(tk.N, tk.S))
        self.listbox.configure(yscrollcommand=scrollbar.set)
        
        # 按钮区域
        btn_frame = ttk.Frame(main_frame, padding="5")
        btn_frame.grid(row=1, column=0, columnspan=2, sticky=(tk.W, tk.E))
        
        ttk.Button(btn_frame, text="选择图片", command=self.select_images).grid(row=0, column=0, padx=5)
        ttk.Button(btn_frame, text="清空列表", command=self.clear_list).grid(row=0, column=1, padx=5)
        
        # 选项区域
        options_frame = ttk.LabelFrame(main_frame, text="选项设置", padding="5")
        options_frame.grid(row=2, column=0, columnspan=2, sticky=(tk.W, tk.E))
        
        # 导出选项
        ttk.Label(options_frame, text="导出方式：").grid(row=0, column=0, padx=5)
        self.export_mode = tk.StringVar(value="single")
        ttk.Radiobutton(options_frame, text="合并为一个文档", variable=self.export_mode, 
                       value="single").grid(row=0, column=1, padx=5)
        ttk.Radiobutton(options_frame, text="每张图片单独导出", variable=self.export_mode, 
                       value="multiple").grid(row=0, column=2, padx=5)
        
        # 语言选项
        ttk.Label(options_frame, text="识别语言：").grid(row=1, column=0, padx=5)
        self.lang_mode = tk.StringVar(value="chi_sim+eng")
        ttk.Radiobutton(options_frame, text="简体中文", variable=self.lang_mode, 
                       value="chi_sim").grid(row=1, column=1, padx=5)
        ttk.Radiobutton(options_frame, text="英文", variable=self.lang_mode, 
                       value="eng").grid(row=1, column=2, padx=5)
        ttk.Radiobutton(options_frame, text="中英双语", variable=self.lang_mode, 
                       value="chi_sim+eng").grid(row=1, column=3, padx=5)
        
        # 导出按钮
        ttk.Button(main_frame, text="开始识别并导出", command=self.process_and_export).grid(row=3, column=0, columnspan=2, pady=10)
        
        # 进度条
        self.progress_var = tk.DoubleVar()
        self.progress = ttk.Progressbar(main_frame, variable=self.progress_var, maximum=100)
        self.progress.grid(row=4, column=0, columnspan=2, sticky=(tk.W, tk.E), pady=5)
        
        # 修改拖放功能的注册方式
        self.listbox.drop_target_register(DND_FILES)
        self.listbox.dnd_bind('<<Drop>>', self.handle_drop)
        
    def select_images(self):
        """选择图片文件"""
        files = filedialog.askopenfilenames(
            title="选择图片",
            filetypes=[("图片文件", "*.png *.jpg *.jpeg")]
        )
        self.add_images(files)
    
    def add_images(self, files):
        """添加图片到列表"""
        for file in files:
            if file not in self.image_paths:
                self.image_paths.append(file)
                self.listbox.insert(tk.END, os.path.basename(file))
    
    def clear_list(self):
        """清空图片列表"""
        self.listbox.delete(0, tk.END)
        self.image_paths.clear()
    
    def handle_drop(self, event):
        """处理拖放事件"""
        # 修改拖放处理方式
        files = self.root.tk.splitlist(event.data)
        valid_files = [f for f in files if f.lower().endswith(('.png', '.jpg', '.jpeg'))]
        if not valid_files:
            messagebox.showwarning("警告", "请只拖放图片文件（PNG/JPG/JPEG）！")
            return
        self.add_images(valid_files)
    
    def process_and_export(self):
        """处理图片并导出文档"""
        if not self.image_paths:
            messagebox.showwarning("警告", "请先添加图片！")
            return
            
        # 选择导出路径
        if self.export_mode.get() == "single":
            save_path = filedialog.asksaveasfilename(
                defaultextension=".docx",
                filetypes=[("Word文档", "*.docx")],
                title="选择保存位置"
            )
            if not save_path:
                return
        else:
            save_dir = filedialog.askdirectory(title="选择保存目录")
            if not save_dir:
                return
        
        try:
            text_results = {}
            total = len(self.image_paths)
            
            for i, image_path in enumerate(self.image_paths):
                processor = ImageProcessor()
                text = processor.recognize_text(image_path, self.lang_mode.get())
                text_results[os.path.basename(image_path)] = text
                
                # 更新进度条
                self.progress_var.set((i + 1) / total * 100)
                self.root.update()
            
            if self.export_mode.get() == "single":
                self.export_single_file(text_results, save_path)
            else:
                self.export_multiple_files(text_results, save_dir)
                
            messagebox.showinfo("成功", "文字识别并导出完成！")
            
        except Exception as e:
            messagebox.showerror("错误", f"处理过程中出现错误：{str(e)}")
        finally:
            self.progress_var.set(0)
    
    def export_single_file(self, text_results, save_path):
        """导出为单个文件"""
        doc = Document()
        for filename, text in text_results.items():
            doc.add_heading(filename, level=1)
            doc.add_paragraph(text)
        doc.save(save_path)
    
    def export_multiple_files(self, text_results, save_dir):
        """导出为多个文件"""
        for filename, text in text_results.items():
            doc = Document()
            doc.add_heading(filename, level=1)
            doc.add_paragraph(text)
            save_path = os.path.join(save_dir, f"{os.path.splitext(filename)[0]}.docx")
            doc.save(save_path) 